package com.example.demo.tokenbucket;

import com.example.demo.redis.RedisUtil;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

@Component
@Slf4j
public class TokenBucketUtil {
    public static final String TOKEN = "token";
    //令牌桶最大数量
    private Integer maxToken = 10;
    //每次放入数量
    private Integer perToken = 2;
    @Resource
    RedisUtil redisUtil;

    @PostConstruct
    @SneakyThrows
    public void init() {
        //程序启动时，初始化令牌数量为最大值
        redisUtil.put(TOKEN, String.valueOf(maxToken));
    }

    @Scheduled(fixedRate = 30000)
    public void putToken() {
        //如果令牌桶未满，且空余数量大于等于perToken = 2 ，则每30秒放入两个令牌
        String LUA_SCRIPT =
                "local currentValue = redis.call('GET', KEYS[1]); " +
                        "if tonumber(currentValue) <= " + (maxToken - perToken) + " then " +
                        "   redis.call('SET', KEYS[1], tonumber(currentValue) + " + perToken + "); " +
                        "else " +
                        "   redis.call('SET', KEYS[1], " + maxToken + "); " +
                        "end " +
                        "return redis.call('GET', KEYS[1]); ";
        Long l = redisUtil.exeLua(LUA_SCRIPT, TOKEN);
        log.info("定时塞入令牌，当前令牌个数：{}", l);
    }
}
